# coding=utf-8
# Copyright 2023-present the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import re
import warnings
from dataclasses import dataclass, field, asdict
from typing import List, Optional, Tuple, Union
from enum import Enum

import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from peft.import_utils import is_bnb_4bit_available, is_bnb_available
from peft.tuners.lora import LoraConfig, LoraModel
from peft.tuners.tuners_utils import BaseTunerLayer
from peft.utils import (
    TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING,
    get_auto_gptq_quant_linear,
    transpose,
    ModulesToSaveWrapper,
    get_quantization_config,
    _get_submodules,
    _freeze_adapter
)
from transformers.pytorch_utils import Conv1D

from .weight_generator import PerceptualWeightsGenerator

if is_bnb_available():
    import bitsandbytes as bnb


@dataclass
class PLoraConfig(LoraConfig):
    """
    This is the configuration class to store the configuration of a [`PLoraModel`].

    Args:
        r (`int`): Lora attention dimension.
        target_modules (`Union[List[str],str]`): The names of the modules to apply Lora to.
        lora_alpha (`int`): The alpha parameter for Lora scaling.
        lora_dropout (`float`): The dropout probability for Lora layers.
        fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out).
            For example, gpt-2 uses `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set
            to `True`.
        bias (`str`): Bias type for Lora. Can be 'none', 'all' or 'lora_only'. If 'all' or 'lora_only', the
            corresponding biases will be updated during training. Be aware that this means that, even when disabling
            the adapters, the model will not produce the same output as the base model would have without adaptation.
        modules_to_save (`List[str]`):List of modules apart from LoRA layers to be set as trainable
            and saved in the final checkpoint.
        layers_to_transform (`Union[List[int],int]`):
            The layer indexes to transform, if this argument is specified, it will apply the LoRA transformations on
            the layer indexes that are specified in this list. If a single integer is passed, it will apply the LoRA
            transformations on the layer at this index.
        layers_pattern (`str`):
            The layer pattern name, used only if `layers_to_transform` is different from `None` and if the layer
            pattern is not in the common layers pattern.
    """
    num_queries: int = field(default=None, metadata={"help": "The number of queries."})
    user_embedding_size: int = field(default=None, metadata={"help": "The embedding size of hidden states."})


class PLoraLayer(BaseTunerLayer):
    def __init__(self, in_features: int, out_features: int, **kwargs):
        self.r = {}
        self.lora_alpha = {}
        self.scaling = {}
        self.lora_dropout = nn.ModuleDict({})
        # self.lora_A = nn.ModuleDict({})
        # self.lora_B = nn.ModuleDict({})
        self.lora_P = nn.ModuleDict({})
        # Mark the weight as unmerged
        self.merged = False
        self.disable_adapters = False
        self.in_features = in_features
        self.out_features = out_features
        self.kwargs = kwargs

    def update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights):
        self.r[adapter_name] = r
        self.lora_alpha[adapter_name] = lora_alpha
        if lora_dropout > 0.0:
            lora_dropout_layer = nn.Dropout(p=lora_dropout)
        else:
            lora_dropout_layer = nn.Identity()

        self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer}))
        # Actual trainable parameters
        if r > 0:
            # self.lora_A.update(nn.ModuleDict({adapter_name: nn.Linear(self.in_features, r, bias=False)}))
            # self.lora_B.update(nn.ModuleDict({adapter_name: nn.Linear(r, self.out_features, bias=False)}))
            self.lora_P.update(nn.ModuleDict({adapter_name: nn.Linear(r, self.out_features, bias=False)}))
            self.scaling[adapter_name] = lora_alpha / r
        if init_lora_weights:
            self.reset_lora_parameters(adapter_name)
        self.to(self.weight.device)

    def reset_lora_parameters(self, adapter_name):
        if adapter_name in self.lora_P.keys():
            # initialize A the same way as the default for nn.Linear and B to zero
            # nn.init.kaiming_uniform_(self.lora_A[adapter_name].weight, a=math.sqrt(5))
            # nn.init.zeros_(self.lora_B[adapter_name].weight)
            nn.init.zeros_(self.lora_P[adapter_name].weight)


# class PLoraModel(torch.nn.Module):
#     def __init__(self, model, config, adapter_name):
#         super().__init__()
#         self.model = model
#         self.forward = self.model.forward
#         self.peft_config = config
#         self.add_adapter(adapter_name, self.peft_config[adapter_name])
#         if self.peft_config[adapter_name].num_queries is not None and self.peft_config[adapter_name].user_embedding_size is not None:
#             self.weight_generator = PerceptualWeightsGenerator(num_queries=config[adapter_name].num_queries,
#                                                                user_embedding_size=config[adapter_name].user_embedding_size,
#                                                                hidden_size=model.config.hidden_size,
#                                                                r=self.peft_config[adapter_name].r)
#         else:
#             self.register_module("weight_generator", None)
#
#     def add_adapter(self, adapter_name, config=None):
#         if config is not None:
#             config = self._prepare_lora_config(config, self.model.config.to_dict())
#             self.peft_config[adapter_name] = config
#         self._find_and_replace(adapter_name)
#         if len(self.peft_config) > 1 and self.peft_config[adapter_name].bias != "none":
#             raise ValueError(
#                 "LoraModel supports only 1 adapter with bias. When using multiple adapters, set bias to 'none' for all adapters."
#             )
#         mark_only_lora_as_trainable(self.model, self.peft_config[adapter_name].bias)
#         if self.peft_config[adapter_name].inference_mode:
#             _freeze_adapter(self.model, adapter_name)
#
#     def _find_and_replace(self, adapter_name):
#         lora_config = self.peft_config[adapter_name]
#         loaded_in_8bit = getattr(self.model, "is_loaded_in_8bit", False)
#         if loaded_in_8bit and not is_bnb_available():
#             raise ImportError(
#                 "To use Lora with 8-bit quantization, please install the `bitsandbytes` package. "
#                 "You can install it with `pip install bitsandbytes`."
#             )
#         is_target_modules_in_base_model = False
#         kwargs = {
#             "r": lora_config.r,
#             "lora_alpha": lora_config.lora_alpha,
#             "lora_dropout": lora_config.lora_dropout,
#             "fan_in_fan_out": lora_config.fan_in_fan_out,
#             "init_lora_weights": lora_config.init_lora_weights,
#         }
#         key_list = [key for key, _ in self.model.named_modules()]
#         for key in key_list:
#             if isinstance(lora_config.target_modules, str):
#                 target_module_found = re.fullmatch(lora_config.target_modules, key)
#             else:
#                 target_module_found = any(key.endswith(target_key) for target_key in lora_config.target_modules)
#             if target_module_found:
#                 if not is_target_modules_in_base_model:
#                     is_target_modules_in_base_model = True
#                 parent, target, target_name = _get_submodules(self.model, key)
#                 bias = target.bias is not None
#                 if isinstance(target, PLoraLayer):
#                     target.update_layer(
#                         adapter_name,
#                         lora_config.r,
#                         lora_config.lora_alpha,
#                         lora_config.lora_dropout,
#                         lora_config.init_lora_weights,
#                     )
#                 else:
#                     if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt):
#                         kwargs.update(
#                             {
#                                 "has_fp16_weights": target.state.has_fp16_weights,
#                                 "memory_efficient_backward": target.state.memory_efficient_backward,
#                                 "threshold": target.state.threshold,
#                                 "index": target.index,
#                             }
#                         )
#                         new_module = Linear8bitLt(
#                             adapter_name, target.in_features, target.out_features, bias=bias, **kwargs
#                         )
#                     else:
#                         if isinstance(target, torch.nn.Linear):
#                             in_features, out_features = target.in_features, target.out_features
#                             if kwargs["fan_in_fan_out"]:
#                                 warnings.warn(
#                                     "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. "
#                                     "Setting fan_in_fan_out to False."
#                                 )
#                                 kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = False
#                         elif isinstance(target, Conv1D):
#                             in_features, out_features = (
#                                 target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape
#                             )
#                             if not kwargs["fan_in_fan_out"]:
#                                 warnings.warn(
#                                     "fan_in_fan_out is set to False but the target module is `Conv1D`. "
#                                     "Setting fan_in_fan_out to True."
#                                 )
#                                 kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = True
#                         else:
#                             raise ValueError(
#                                 f"Target module {target} is not supported. "
#                                 f"Currently, only `torch.nn.Linear` and `Conv1D` are supported."
#                             )
#                         new_module = Linear(adapter_name, in_features, out_features, bias=bias, **kwargs)
#
#                     self._replace_module(parent, target_name, new_module, target)
#         if not is_target_modules_in_base_model:
#             raise ValueError(
#                 f"Target modules {lora_config.target_modules} not found in the base model. "
#                 f"Please check the target modules and try again."
#             )
#
#     def _replace_module(self, parent_module, child_name, new_module, old_module):
#         setattr(parent_module, child_name, new_module)
#         new_module.weight = old_module.weight
#         if old_module.bias is not None:
#             new_module.bias = old_module.bias
#         if getattr(old_module, "state", None) is not None:
#             new_module.state = old_module.state
#             new_module.to(old_module.weight.device)
#
#         # dispatch to correct device
#         for name, module in new_module.named_modules():
#             if "lora_" in name:
#                 module.to(old_module.weight.device)
#
#     def __getattr__(self, name: str):
#         """Forward missing attributes to the wrapped module."""
#         try:
#             return super().__getattr__(name)  # defer to nn.Module's logic
#         except AttributeError:
#             return getattr(self.model, name)
#
#     @property
#     def modules_to_save(self):
#         return None
#
#     def get_peft_config_as_dict(self, inference: bool = False):
#         config_dict = {}
#         for key, value in self.peft_config.items():
#             config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()}
#             if inference:
#                 config["inference_mode"] = True
#         config_dict[key] = config
#         return config
#
#     def _set_adapter_layers(self, enabled=True):
#         for module in self.model.modules():
#             if isinstance(module, PLoraLayer):
#                 module.disable_adapters = False if enabled else True
#
#     def enable_adapter_layers(self):
#         self._set_adapter_layers(enabled=True)
#
#     def disable_adapter_layers(self):
#         self._set_adapter_layers(enabled=False)
#
#     def set_adapter(self, adapter_name):
#         for module in self.model.modules():
#             if isinstance(module, PLoraLayer):
#                 if module.merged:
#                     warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.")
#                     module.unmerge()
#                 module.active_adapter = adapter_name
#
#     def merge_adapter(self):
#         for module in self.model.modules():
#             if isinstance(module, PLoraLayer):
#                 module.merge()
#
#     def unmerge_adapter(self):
#         for module in self.model.modules():
#             if isinstance(module, PLoraLayer):
#                 module.unmerge()
#
#     @staticmethod
#     def _prepare_lora_config(peft_config, model_config):
#         if peft_config.target_modules is None:
#             if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
#                 raise ValueError("Please specify `target_modules` in `peft_config`")
#             peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]]
#         if peft_config.inference_mode:
#             peft_config.merge_weights = True
#         return peft_config
#
#     def merge_and_unload(self):
#         r"""
#         This method merges the LoRa layers into the base model. This is needed if someone wants to use the base model
#         as a standalone model.
#         """
#         if self.config.model_type == "gpt2":
#             raise ValueError("GPT2 models are not supported for merging LORA layers")
#
#         if getattr(self.model, "is_loaded_in_8bit", False):
#             raise ValueError("Cannot merge LORA layers when the model is loaded in 8-bit mode")
#
#         key_list = [key for key, _ in self.model.named_modules() if "lora" not in key]
#         for key in key_list:
#             parent, target, target_name = _get_submodules(self.model, key)
#             if isinstance(target, PLoraLayer):
#                 bias = target.bias is not None
#                 new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias)
#                 target.merge()
#                 self._replace_module(parent, target_name, new_module, target)
#         return self.model


def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None:
    for n, p in model.named_parameters():
        if "lora_" not in n:
            p.requires_grad = False
    if bias == "none":
        return
    elif bias == "all":
        for n, p in model.named_parameters():
            if "bias" in n:
                p.requires_grad = True
    elif bias == "lora_only":
        for m in model.modules():
            if isinstance(m, PLoraLayer) and hasattr(m, "bias") and m.bias is not None:
                m.bias.requires_grad = True
    else:
        raise NotImplementedError


class PLoraModel(LoraModel):
    """
    Creates Low Rank Adapter (Lora) model from a pretrained transformers model.

    Args:
        model ([`~transformers.PreTrainedModel`]): The model to be adapted.
        config ([`PLoraConfig`]): The configuration of the Lora model.
        adapter_name (`str`): The name of the adapter, defaults to `"default"`.

    Returns:
        `torch.nn.Module`: The Lora model.
        ```
    **Attributes**:
        - **model** ([`~transformers.PreTrainedModel`]) -- The model to be adapted.
        - **peft_config** ([`PLoraConfig`]): The configuration of the Lora model.
    """

    def __init__(self, model, config, adapter_name) -> None:
        super().__init__(model, config, adapter_name)
        # self.forward = self.model.forward
        if self.peft_config[adapter_name].num_queries is not None and self.peft_config[adapter_name].user_embedding_size is not None:
            self.weight_generator = PerceptualWeightsGenerator(num_queries=config[adapter_name].num_queries,
                                                               user_embedding_size=config[adapter_name].user_embedding_size,
                                                               hidden_size=model.config.hidden_size,
                                                               r=self.peft_config[adapter_name].r)
        else:
            self.register_module("weight_generator", None)

    def _create_and_replace(
        self,
        lora_config,
        adapter_name,
        target,
        target_name,
        parent,
        **optionnal_kwargs,
    ):
        bias = hasattr(target, "bias") and target.bias is not None
        kwargs = {
            "r": lora_config.r,
            "lora_alpha": lora_config.lora_alpha,
            "lora_dropout": lora_config.lora_dropout,
            "fan_in_fan_out": lora_config.fan_in_fan_out,
            "init_lora_weights": lora_config.init_lora_weights,
            "loaded_in_8bit": optionnal_kwargs.pop("loaded_in_8bit", False),
            "loaded_in_4bit": optionnal_kwargs.pop("loaded_in_4bit", False),
            "bias": bias
        }

        quantization_config = get_quantization_config(self.model, method="gptq")
        if quantization_config is not None:
            kwargs["gptq_quantization_config"] = quantization_config

        # TODO: better deal with that
        if isinstance(target, PLoraLayer) and isinstance(target, torch.nn.Conv2d):
            target.update_layer_conv2d(
                adapter_name,
                lora_config.r,
                lora_config.lora_alpha,
                lora_config.lora_dropout,
                lora_config.init_lora_weights,
            )
        elif isinstance(target, PLoraLayer) and isinstance(target, torch.nn.Embedding):
            target.update_layer_embedding(
                adapter_name,
                lora_config.r,
                lora_config.lora_alpha,
                lora_config.lora_dropout,
                lora_config.init_lora_weights,
            )

        elif isinstance(target, PLoraLayer):
            target.update_layer(
                adapter_name,
                lora_config.r,
                lora_config.lora_alpha,
                lora_config.lora_dropout,
                lora_config.init_lora_weights,
            )
        else:
            new_module = self._create_new_module(lora_config, adapter_name, target, **kwargs)
            self._replace_module(parent, target_name, new_module, target)

    @staticmethod
    def _replace_module(parent, child_name, new_module, child):
        setattr(parent, child_name, new_module)
        new_module.weight = child.weight
        if hasattr(child, "bias"):
            if child.bias is not None:
                new_module.bias = child.bias

        if getattr(child, "state", None) is not None:
            new_module.state = child.state
            new_module.to(child.weight.device)

        # dispatch to correct device
        for name, module in new_module.named_modules():
            if "lora_" in name:
                module.to(child.weight.device)
            if "ranknum" in name:
                module.to(child.weight.device)

    def _mark_only_adapters_as_trainable(self) -> None:
        active_adapter = self._get_active_adapter()
        bias = self.peft_config[active_adapter].bias

        for n, p in self.model.named_parameters():
            if "lora_" not in n:
                p.requires_grad = False
        if bias == "none":
            return
        elif bias == "all":
            for n, p in self.model.named_parameters():
                if "bias" in n:
                    p.requires_grad = True
        elif bias == "lora_only":
            for m in self.model.modules():
                if isinstance(m, PLoraLayer) and hasattr(m, "bias") and m.bias is not None:
                    m.bias.requires_grad = True
        else:
            raise NotImplementedError

    @staticmethod
    def _create_new_module(lora_config, adapter_name, target, **kwargs):
        gptq_quantization_config = kwargs.get("gptq_quantization_config", None)
        AutoGPTQQuantLinear = get_auto_gptq_quant_linear(gptq_quantization_config)

        loaded_in_8bit = kwargs.pop("loaded_in_8bit", False)
        loaded_in_4bit = kwargs.pop("loaded_in_4bit", False)
        bias = kwargs.pop("bias", False)

        if loaded_in_8bit and isinstance(target, bnb.nn.Linear8bitLt):
            eightbit_kwargs = kwargs.copy()
            eightbit_kwargs.update(
                {
                    "has_fp16_weights": target.state.has_fp16_weights,
                    "memory_efficient_backward": target.state.memory_efficient_backward,
                    "threshold": target.state.threshold,
                    "index": target.index,
                }
            )
            new_module = Linear8bitLt(
                adapter_name, target.in_features, target.out_features, bias=bias, **eightbit_kwargs
            )
        elif loaded_in_4bit and is_bnb_4bit_available() and isinstance(target, bnb.nn.Linear4bit):
            fourbit_kwargs = kwargs.copy()
            fourbit_kwargs.update(
                {
                    "compute_dtype": target.compute_dtype,
                    "compress_statistics": target.weight.compress_statistics,
                    "quant_type": target.weight.quant_type,
                }
            )
            new_module = Linear4bit(adapter_name, target.in_features, target.out_features, bias=bias, **fourbit_kwargs)
        elif AutoGPTQQuantLinear is not None and isinstance(target, AutoGPTQQuantLinear):
            new_module = QuantLinear(adapter_name, target, **kwargs)
            target.weight = target.qweight
        elif isinstance(target, torch.nn.Embedding):
            embedding_kwargs = kwargs.copy()
            embedding_kwargs.pop("fan_in_fan_out", None)
            in_features, out_features = target.num_embeddings, target.embedding_dim
            new_module = Embedding(adapter_name, in_features, out_features, **embedding_kwargs)
        elif isinstance(target, torch.nn.Conv2d):
            out_channels, in_channels = target.weight.size()[:2]
            kernel_size = target.weight.size()[2:]
            stride = target.stride
            padding = target.padding
            new_module = Conv2d(adapter_name, in_channels, out_channels, kernel_size, stride, padding, **kwargs)
        else:
            if isinstance(target, torch.nn.Linear):
                in_features, out_features = target.in_features, target.out_features
                if kwargs["fan_in_fan_out"]:
                    warnings.warn(
                        "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. "
                        "Setting fan_in_fan_out to False."
                    )
                    kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = False
            elif isinstance(target, Conv1D):
                in_features, out_features = (
                    target.weight.ds_shape if hasattr(target.weight, "ds_shape") else target.weight.shape
                )
                kwargs["is_target_conv_1d_layer"] = True
                if not kwargs["fan_in_fan_out"]:
                    warnings.warn(
                        "fan_in_fan_out is set to False but the target module is `Conv1D`. "
                        "Setting fan_in_fan_out to True."
                    )
                    kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = True
            else:
                raise ValueError(
                    f"Target module {target} is not supported. "
                    f"Currently, only `torch.nn.Linear` and `Conv1D` are supported."
                )
            new_module = Linear(adapter_name, in_features, out_features, bias=bias, **kwargs)

        return new_module

    def _set_adapter_layers(self, enabled=True):
        for module in self.model.modules():
            if isinstance(module, PLoraLayer):
                module.disable_adapters = False if enabled else True
            elif isinstance(module, ModulesToSaveWrapper):
                module.disable_adapters = False if enabled else True

    def enable_adapter_layers(self):
        self._set_adapter_layers(enabled=True)

    def _get_active_adapter(self) -> str:
        active_adapter = None
        for module in self.model.modules():
            if isinstance(module, PLoraLayer):
                active_adapter = module.active_adapter

        if active_adapter is None:
            raise ValueError(
                "Something went wrong, no active adapter could be found, please report the issue on GitHub"
            )
        return active_adapter

    def disable_adapter_layers(self):
        active_adapter = self._get_active_adapter()
        val = self.peft_config[active_adapter].bias
        if val != "none":
            msg = (
                f"Careful, disabling adapter layers with bias configured to be '{val}' does not produce the same "
                "output as the the base model would without adaption."
            )
            warnings.warn(msg)
        self._set_adapter_layers(enabled=False)

    def set_adapter(self, adapter_name):
        for module in self.model.modules():
            if isinstance(module, PLoraLayer):
                if module.merged:
                    warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.")
                    module.unmerge()
                module.active_adapter = adapter_name

    def _unload_and_optionally_merge(self, merge=True, progressbar: bool = False):
        if merge:
            if getattr(self.model, "is_loaded_in_8bit", False) or getattr(self.model, "is_loaded_in_4bit", False):
                raise ValueError("Cannot merge LORA layers when the model is loaded in 8-bit mode")
            if getattr(self.model, "quantization_method", None) == "gptq":
                raise ValueError("Cannot merge LORA layers when the model is gptq quantized")

        key_list = [key for key, _ in self.model.named_modules() if "lora" not in key]
        desc = "Unloading " + ("and merging " if merge else "") + "model"
        for key in tqdm(key_list, disable=not progressbar, desc=desc):
            try:
                parent, target, target_name = _get_submodules(self.model, key)
            except AttributeError:
                continue
            if isinstance(target, PLoraLayer):
                if isinstance(target, nn.Embedding):
                    new_module = torch.nn.Embedding(target.in_features, target.out_features)
                elif isinstance(target, nn.Conv2d):
                    new_module = torch.nn.Conv2d(
                        target.in_channels,
                        target.out_channels,
                        kernel_size=target.kernel_size,
                        stride=target.stride,
                        padding=target.padding,
                        dilation=target.dilation,
                    )
                else:
                    bias = target.bias is not None
                    if getattr(target, "is_target_conv_1d_layer", False):
                        new_module = Conv1D(target.out_features, target.in_features)
                    else:
                        new_module = torch.nn.Linear(target.in_features, target.out_features, bias=bias)
                if merge:
                    target.merge()
                self._replace_module(parent, target_name, new_module, target)

            # save any additional trainable modules part of `modules_to_save`
            if isinstance(target, ModulesToSaveWrapper):
                setattr(parent, target_name, target.modules_to_save[target.active_adapter])

        return self.model


class Linear(nn.Linear, PLoraLayer):
    # Lora implemented in a dense layer
    def __init__(
            self,
            adapter_name: str,
            in_features: int,
            out_features: int,
            r: int = 0,
            lora_alpha: int = 1,
            lora_dropout: float = 0.0,
            fan_in_fan_out: bool = False,
            # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
            is_target_conv_1d_layer: bool = False,
            **kwargs,
    ):
        init_lora_weights = kwargs.pop("init_lora_weights", True)

        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        PLoraLayer.__init__(self, in_features=in_features, out_features=out_features)
        # Freezing the pre-trained weight matrix
        self.weight.requires_grad = False

        self.fan_in_fan_out = fan_in_fan_out
        if fan_in_fan_out:
            self.weight.data = self.weight.data.T

        nn.Linear.reset_parameters(self)
        self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
        self.active_adapter = adapter_name
        self.is_target_conv_1d_layer = is_target_conv_1d_layer

    def merge(self):
        if self.active_adapter not in self.lora_P.keys():
            return
        if self.merged:
            warnings.warn("Already merged. Nothing to do.")
            return
        if self.r[self.active_adapter] > 0:
            self.weight.data += self.get_delta_weight(self.active_adapter)
            self.merged = True

    def unmerge(self):
        if self.active_adapter not in self.lora_P.keys():
            return
        if not self.merged:
            warnings.warn("Already unmerged. Nothing to do.")
            return
        if self.r[self.active_adapter] > 0:
            self.weight.data -= self.get_delta_weight(self.active_adapter)
            self.merged = False

    def get_delta_weight(self, adapter):
        return (
                transpose(
                    self.lora_B[adapter].weight @ self.lora_A[adapter].weight,
                    self.fan_in_fan_out,
                )
                * self.scaling[adapter]
        )

    def forward(self, x: torch.Tensor, w: Optional[torch.Tensor] = None):
        weight = w
        user, item = torch.split(w, w.shape[-2]//2, dim=-2)
        previous_dtype = x.dtype
        if self.active_adapter not in self.lora_P.keys():
            return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
        if self.disable_adapters:
            if self.r[self.active_adapter] > 0 and self.merged:
                self.unmerge()
            result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
        elif self.r[self.active_adapter] > 0 and not self.merged:
            result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)

            x = x.to(self.lora_P[self.active_adapter].weight.dtype)

            if weight is not None:
                result += (
                        self.lora_P[self.active_adapter](
                            torch.bmm(self.lora_dropout[self.active_adapter](x), weight)
                        )
                        * self.scaling[self.active_adapter]
                )
        else:
            result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)

        result = result.to(previous_dtype)

        return result


class Embedding(nn.Embedding, PLoraLayer):
    # LoRA implemented in a Embedding layer
    def __init__(
            self,
            adapter_name: str,
            num_embeddings: int,
            embedding_dim: int,
            r: int = 0,
            lora_alpha: int = 1,
            lora_dropout: float = 0.0,
            **kwargs,
    ):
        init_lora_weights = kwargs.pop("init_lora_weights", True)

        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
        PLoraLayer.__init__(self, in_features=num_embeddings, out_features=embedding_dim)

        self.weight.requires_grad = False

        nn.Embedding.reset_parameters(self)
        self.update_layer_embedding(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
        self.active_adapter = adapter_name

    def unmerge(self):
        if not self.merged:
            warnings.warn("Already unmerged. Nothing to do.")
            return
        if self.r[self.active_adapter] > 0:
            self.weight.data -= self.get_delta_weight(self.active_adapter)
            self.merged = False

    def merge(self):
        if self.merged:
            warnings.warn("Already merged. Nothing to do.")
            return
        if self.r[self.active_adapter] > 0:
            self.weight.data += self.get_delta_weight(self.active_adapter)
            self.merged = True

    def get_delta_weight(self, adapter):
        return transpose(self.lora_embedding_B[adapter] @ self.lora_embedding_A[adapter], True) * self.scaling[adapter]

    def forward(self, x: torch.Tensor):
        if self.disable_adapters:
            if self.r[self.active_adapter] > 0 and self.merged:
                self.unmerge()
            return nn.Embedding.forward(self, x)

        elif self.r[self.active_adapter] > 0 and not self.merged:
            result = nn.Embedding.forward(self, x)
            if self.r[self.active_adapter] > 0:
                after_A = F.embedding(
                    x,
                    self.lora_embedding_A[self.active_adapter].T,
                    self.padding_idx,
                    self.max_norm,
                    self.norm_type,
                    self.scale_grad_by_freq,
                    self.sparse,
                )
                result += (after_A @ self.lora_embedding_B[self.active_adapter].T) * self.scaling[self.active_adapter]
            return result
        else:
            return nn.Embedding.forward(self, x)


class Conv2d(nn.Conv2d, PLoraLayer):
    # Lora implemented in a conv2d layer
    def __init__(
            self,
            adapter_name: str,
            in_channels: int,
            out_channels: int,
            kernel_size: Union[int, Tuple[int]],
            stride: Union[int, Tuple[int]] = 1,
            padding: Union[int, Tuple[int]] = 0,
            r: int = 0,
            lora_alpha: int = 1,
            lora_dropout: float = 0.0,
            **kwargs,
    ):
        init_lora_weights = kwargs.pop("init_lora_weights", True)

        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride, padding)
        PLoraLayer.__init__(
            self,
            in_features=in_channels,
            out_features=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
        )
        # Freezing the pre-trained weight matrix
        self.weight.requires_grad = False

        nn.Conv2d.reset_parameters(self)
        self.update_layer_conv2d(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
        self.active_adapter = adapter_name

    def merge(self):
        if self.active_adapter not in self.lora_A.keys():
            return
        if self.merged:
            warnings.warn("Already merged. Nothing to do.")
            return
        if self.r[self.active_adapter] > 0:
            self.weight.data += self.get_delta_weight(self.active_adapter)
            self.merged = True

    def unmerge(self):
        if self.active_adapter not in self.lora_A.keys():
            return
        if not self.merged:
            warnings.warn("Already unmerged. Nothing to do.")
            return
        if self.r[self.active_adapter] > 0:
            self.weight.data -= self.get_delta_weight(self.active_adapter)
            self.merged = False

    def get_delta_weight(self, adapter):
        # https://github.com/bmaltais/kohya_ss/blob/feb6728762a8f463d15ba936d189d4c3abfaa1ab/networks/lora.py#L117
        if self.weight.size()[2:4] == (1, 1):
            # conv2d 1x1
            return (
                    self.lora_B[adapter].weight.squeeze(3).squeeze(2) @ self.lora_A[adapter].weight.squeeze(3).squeeze(
                2)
            ).unsqueeze(2).unsqueeze(3) * self.scaling[adapter]
        else:
            # conv2d 3x3
            return (
                    F.conv2d(
                        self.lora_A[adapter].weight.permute(1, 0, 2, 3),
                        self.lora_B[adapter].weight,
                    ).permute(1, 0, 2, 3)
                    * self.scaling[adapter]
            )

    def forward(self, x: torch.Tensor):
        previous_dtype = x.dtype

        if self.active_adapter not in self.lora_A.keys():
            return F.conv2d(
                x,
                self.weight,
                bias=self.bias,
                stride=self.stride,
                padding=self.padding,
                dilation=self.dilation,
                groups=self.groups,
            )
        if self.disable_adapters:
            if self.r[self.active_adapter] > 0 and self.merged:
                self.unmerge()
            result = F.conv2d(
                x,
                self.weight,
                bias=self.bias,
                stride=self.stride,
                padding=self.padding,
                dilation=self.dilation,
                groups=self.groups,
            )
        elif self.r[self.active_adapter] > 0 and not self.merged:
            result = F.conv2d(
                x,
                self.weight,
                bias=self.bias,
                stride=self.stride,
                padding=self.padding,
                dilation=self.dilation,
                groups=self.groups,
            )

            x = x.to(self.lora_A[self.active_adapter].weight.dtype)

            result += (
                    self.lora_B[self.active_adapter](
                        self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
                    )
                    * self.scaling[self.active_adapter]
            )
        else:
            result = F.conv2d(
                x,
                self.weight,
                bias=self.bias,
                stride=self.stride,
                padding=self.padding,
                dilation=self.dilation,
                groups=self.groups,
            )

        result = result.to(previous_dtype)

        return result


if is_bnb_available():

    class Linear8bitLt(bnb.nn.Linear8bitLt, PLoraLayer):
        # Lora implemented in a dense layer
        def __init__(
                self,
                adapter_name,
                in_features,
                out_features,
                r: int = 0,
                lora_alpha: int = 1,
                lora_dropout: float = 0.0,
                **kwargs,
        ):
            bnb.nn.Linear8bitLt.__init__(
                self,
                in_features,
                out_features,
                bias=kwargs.get("bias", True),
                has_fp16_weights=kwargs.get("has_fp16_weights", True),
                memory_efficient_backward=kwargs.get("memory_efficient_backward", False),
                threshold=kwargs.get("threshold", 0.0),
                index=kwargs.get("index", None),
            )
            PLoraLayer.__init__(self, in_features=in_features, out_features=out_features)

            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
            init_lora_weights = kwargs.pop("init_lora_weights", True)
            self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
            self.active_adapter = adapter_name

        def forward(self, x: torch.Tensor):
            result = super().forward(x)

            if self.disable_adapters or self.active_adapter not in self.lora_A.keys():
                return result
            elif self.r[self.active_adapter] > 0:
                if not torch.is_autocast_enabled():
                    expected_dtype = result.dtype

                    if x.dtype != torch.float32:
                        x = x.float()
                    output = (
                            self.lora_B[self.active_adapter](
                                self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
                            ).to(expected_dtype)
                            * self.scaling[self.active_adapter]
                    )
                else:
                    output = (
                            self.lora_B[self.active_adapter](
                                self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
                            )
                            * self.scaling[self.active_adapter]
                    )
                result += output
            return result


    if is_bnb_4bit_available():

        class Linear4bit(bnb.nn.Linear4bit, PLoraLayer):
            # Lora implemented in a dense layer
            def __init__(
                    self,
                    adapter_name,
                    in_features,
                    out_features,
                    r: int = 0,
                    lora_alpha: int = 1,
                    lora_dropout: float = 0.0,
                    **kwargs,
            ):
                bnb.nn.Linear4bit.__init__(
                    self,
                    in_features,
                    out_features,
                    bias=kwargs.get("bias", True),
                    compute_dtype=kwargs.get("compute_dtype", torch.float32),
                    compress_statistics=kwargs.get("compress_statistics", True),
                    quant_type=kwargs.get("quant_type", "nf4"),
                )
                PLoraLayer.__init__(self, in_features=in_features, out_features=out_features)

                # Freezing the pre-trained weight matrix
                self.weight.requires_grad = False

                init_lora_weights = kwargs.pop("init_lora_weights", True)
                self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
                self.active_adapter = adapter_name

            def forward(self, x: torch.Tensor):
                result = super().forward(x)

                if self.disable_adapters or self.active_adapter not in self.lora_A.keys():
                    return result
                elif self.r[self.active_adapter] > 0:
                    result = result.clone()
                    if not torch.is_autocast_enabled():
                        expected_dtype = result.dtype
                        x = x.to(self.lora_A[self.active_adapter].weight.dtype)
                        output = (
                                self.lora_B[self.active_adapter](
                                    self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
                                ).to(expected_dtype)
                                * self.scaling[self.active_adapter]
                        )
                    else:
                        output = (
                                self.lora_B[self.active_adapter](
                                    self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
                                )
                                * self.scaling[self.active_adapter]
                        )
                    result += output
                return result


class QuantLinear(torch.nn.Module, PLoraLayer):
    def __init__(
            self,
            adapter_name,
            quant_linear_module,
            r: int = 0,
            lora_alpha: int = 1,
            lora_dropout: float = 0.0,
            **kwargs,
    ):
        torch.nn.Module.__init__(self)
        PLoraLayer.__init__(
            self, in_features=quant_linear_module.infeatures, out_features=quant_linear_module.outfeatures
        )
        self.quant_linear_module = quant_linear_module
        self.weight = quant_linear_module.qweight
        init_lora_weights = kwargs.pop("init_lora_weights", True)
        self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights)
        self.active_adapter = adapter_name

    def forward(self, x: torch.Tensor):
        result = self.quant_linear_module(x)
        if self.disable_adapters or self.active_adapter not in self.lora_A.keys():
            return result
        elif self.r[self.active_adapter] > 0:
            result = result.clone()
            if not torch.is_autocast_enabled():
                expected_dtype = result.dtype
                x = x.to(self.lora_A[self.active_adapter].weight.dtype)
                output = (
                        self.lora_B[self.active_adapter](
                            self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
                        ).to(expected_dtype)
                        * self.scaling[self.active_adapter]
                )
            else:
                output = (
                        self.lora_B[self.active_adapter](
                            self.lora_A[self.active_adapter](self.lora_dropout[self.active_adapter](x))
                        )
                        * self.scaling[self.active_adapter]
                )
            result += output
        return result

    # TODO: Check if it is better as suggested by users https://github.com/PanQiWei/AutoGPTQ/pull/102
    # def reset_lora_parameters(self, adapter_name):
    #     if adapter_name in self.lora_A.keys():
    #         torch.nn.init.xavier_uniform_(self.lora_A[adapter_name].weight)
    #         torch.nn.init.zeros_(self.lora_B[adapter_name].weight)
